import numpy as np
import cv2 as cv

debug = True


class Camera(cv.VideoCapture):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.__is_open = False
        self.__is_in_calibrated = False
        self.__is_ex_calibrated = False
        self.__is_calibrated = False
        self.in_mtx = np.zeros((3, 3))
        self.dist = np.zeros((1, 5))
        self.ex_mtx = np.zeros((3, 4))
        self.p2w_mtx = np.zeros((3, 3))

    def calibrate_intrinsic_matrix(self, images, col, row):
        """
        Calibrate intrinsic matrix of camera
        :param images:
        :param col:
        :param row:
        :return: True if calibration is successful, False otherwise.
        """
        col = int(col)
        row = int(row)

        # termination criteria
        criteria = (cv.TERM_CRITERIA_MAX_ITER + cv.TERM_CRITERIA_EPS, 100, 1e-4)

        # prepare object points, like (0,0,0), (1,0,0), (2,0,0) ....,(6,5,0)
        obj_p = np.zeros((col * row, 3), np.float32)
        obj_p[:, :2] = np.mgrid[0:col, 0:row].T.reshape(-1, 2)
        # Arrays to store object points and image points from all the images.
        obj_points = []  # 3d point in real world space
        img_points = []  # 2d points in image plane

        img_size = 0

        for img in images:
            gray = cv.cvtColor(img, cv.COLOR_BGR2GRAY)
            img_size = gray.shape[::-1]

            # Find the chess board corners
            ret, corners = cv.findChessboardCorners(gray, (col, row), None)

            # If found, add object points, image points (after refining them)
            if ret:
                obj_points.append(obj_p)

                corners2 = cv.cornerSubPix(gray, corners, (11, 11), (-1, -1), criteria)
                img_points.append(corners2)

                # if debug:
                #     # Draw and display the corners
                #     print(corners)
                #     cv.drawChessboardCorners(img, (col, row), corners, ret)
                #     cv.imshow('img', img)
                #     cv.waitKey(0)
                #     cv.destroyWindow('img')

        # Do calibration
        ret, mtx, dist, _, _ = cv.calibrateCamera(obj_points, img_points, img_size, None, None)

        if ret:
            self.in_mtx = mtx
            self.dist = dist
            self.__is_in_calibrated = True
            return True
        else:
            return False

    def calibrate_extrinsic_matrix(self, img_points, obj_points):
        """
        Calibrate extrinsic matrix of camera
        :param img_points: image points coordinates in adjusted images
        :param obj_points:
        :return: True if calibration is successful, False otherwise
        """
        # Intrinsic matrix must be calibrated first
        if not self.__is_in_calibrated:
            return False

        # distortion = np.zeros((1, 5))
        obj_points = np.array(obj_points)
        img_points = np.array(img_points)
        ret, r_vec, t_vec = cv.solvePnP(obj_points, img_points, self.in_mtx, self.dist)
        if not ret:
            return False

        r_mat, _ = cv.Rodrigues(r_vec)
        self.ex_mtx = np.c_[r_mat, t_vec]
        self.__is_ex_calibrated = True

        world2pixel_mtx = self.in_mtx.dot(self.ex_mtx)
        world2pixel_mtx3 = world2pixel_mtx[:, (0, 1, 3)]
        world2pixel_mtx3 = np.matrix(world2pixel_mtx3)
        self.p2w_mtx = world2pixel_mtx3.I
        self.__is_calibrated = True

        return True

    def adjust_image(self, img):
        """
        Adjust image according to intrinsic matrix
        :param img:
        :return: Adjusted image, False if intrinsic matrix is not calibrated
        """
        if not self.__is_in_calibrated:
            return False

        h, w = img.shape[:2]
        new_mtx, roi = cv.getOptimalNewCameraMatrix(self.in_mtx, self.dist, (w, h), 1, (w, h))

        # Undistort
        dst = cv.undistort(img, self.in_mtx, self.dist, None, new_mtx)

        # crop the image
        x, y, w, h = roi
        dst = dst[y:y + h, x:x + w]
        return dst

    def find_objects(self, color):
        """
        Find center coordinates of all objects of color
        :param color: color can be 'green', 'blue' or 'pink'
        :return: List of all center coordinates, False if camera is not calibrated
        """
        # if not self.__is_calibrated:
        #     return False

        if color == 'green':
            lower = (45, 100, 120)
            upper = (90, 255, 255)
        elif color == 'blue':
            lower = (125, 100, 120)
            upper = (160, 255, 255)
        elif color == 'pink':
            lower = (220, 100, 120)
            upper = (245, 255, 255)
        else:
            return False

        # read image
        ret, img = self.read()
        if debug:
            cv.imshow('debug', img)
            cv.waitKey(500)
            cv.destroyWindow('debug')
        # img = self.adjust_image(img)
        img = cv.GaussianBlur(img, (5, 5), 0)

        img_hsv = cv.cvtColor(img, cv.COLOR_BGR2HSV_FULL)

        # get color range
        mask = cv.inRange(img_hsv, lower, upper)
        # if debug:
        #     cv.imshow('test', mask)
        #     cv.waitKey(50)
        #     cv.destroyWindow('test')

        kernel = (177, 177)

        # erode
        mask = cv.erode(mask, kernel, iterations=7)

        # morphology
        mask = cv.morphologyEx(mask, cv.MORPH_OPEN, kernel, iterations=5)

        # dilate
        mask = cv.dilate(mask, kernel, iterations=7)
        if debug:
            cv.imshow('dilate', mask)
            cv.waitKey(500)
            cv.destroyWindow('dilate')

        # get objects contours
        contours, heir = cv.findContours(mask, cv.RETR_EXTERNAL, cv.CHAIN_APPROX_NONE)

        obj_points = []
        for contour in contours:
            img_point = contour.mean(axis=0)
            print(img_point)
            img_point = np.c_[img_point, np.array(1)].T

            center = (int(img_point[0, 0]), int(img_point[1, 0]))
            cv.line(img, center, center, (0, 0, 225), 10)
            # obj_points.append(center)

            point_vec = self.p2w_mtx.dot(img_point)
            obj_point = np.array([point_vec[0, 0] / point_vec[2, 0], point_vec[1, 0] / point_vec[2, 0]])
            obj_points.append(obj_point)

        cv.imshow('points', img)
        cv.waitKey(0)
        cv.destroyWindow('points')

        return obj_points

    def find_flag_point(self):
        # if not self.__is_in_calibrated:
        #     return False

        lower = (15, 130, 160)
        upper = (35, 255, 255)

        # read image
        ret, img = self.read()
        # img = self.adjust_image(img)
        img = cv.GaussianBlur(img, (3, 3), 0)
        img_hsv = cv.cvtColor(img, cv.COLOR_BGR2HSV_FULL)

        # get color range
        mask = cv.inRange(img_hsv, lower, upper)
        # if debug:
        #     cv.imshow('test', mask)
        #     cv.waitKey(0)
        #     cv.destroyWindow('test')

        kernel = (237, 237)

        # erode
        mask = cv.erode(mask, kernel, iterations=7)
        # cv.imshow('debug', mask)
        # cv.waitKey(0)

        # morphology
        mask = cv.morphologyEx(mask, cv.MORPH_OPEN, kernel, iterations=5)

        # dilate
        mask = cv.dilate(mask, kernel, iterations=7)
        # cv.imshow('debug', mask)
        # cv.waitKey(0)

        # get objects contours
        contours, heir = cv.findContours(mask, cv.RETR_EXTERNAL, cv.CHAIN_APPROX_NONE)

        contour = contours[0]
        img_point = contour.mean(axis=0)
        if debug:
            center = (int(img_point[0, 0]), int(img_point[0, 1]))
            cv.line(img, center, center, (0, 0, 225), 10)

            cv.imshow('debug', img)
            cv.waitKey(500)
            cv.destroyWindow('debug')

        return img_point[0]


